import numpy as np
import pymc as pm
import matplotlib.pyplot as plt
from typing import Tuple, List, Dict, Any
import arviz as az
import time
import warnings
warnings.filterwarnings('ignore')

from TDS import TDS
from scipy import interpolate
from dispCal.disp import calDisp

# 定义黑箱似然函数的PyTensor操作
class DispersionLikelihoodOp(Op):
    """
    黑箱似然函数的PyTensor操作类
    """
    
    def __init__(self, freqs, vels, z, 
                 obsed_FV,
                 f_control,v_control,DEBUG=False):
        self.freqs = freqs
        self.vels = vels
        self.z = z
        self.obsed_FV = obsed_FV
        self.V, self.F = np.meshgrid(vels, freqs)
        self.n_layers = len(z)
        self.DEBUG = DEBUG
        # 初始化TDS对象
        # 固定参数
        self.vp = np.array([1.5] * self.n_layers)
        self.qs = np.array([1000] * self.n_layers)
        self.qp = np.array([1000] * self.n_layers)
        self.rho = np.array([2.0] * self.n_layers)
        self.z_extended = np.concatenate([z, [z[-1] + 5.0]])

        self.thickness = np.diff(self.z_extended)
        self.thickness[-1]=0.0
        self.f_control = f_control
        self.v_control = v_control

        self.METHOD='curve' # or spec
        self.sigma=0.05
    
    def get_spectrum(self, vs):

        # 设置媒质参数
        ts = time.time()
        tds = TDS(DEBUG=self.DEBUG)
        tds.set_media_parameters(
            z=self.z_extended, 
            vp0=self.vp, 
            vs0=vs, 
            rho=self.rho, 
            qs=self.qs, 
            qp=self.qp
        )

        # 设置源参数和计算参数
        tds.set_source_parameters(
            zs=0.05, z0=0.0, 
            mino=0.1, maxo=4.0, no=len(self.freqs),
            minv=np.min(vs)/2, maxv=2, nv=len(self.vels)
        )

        tds.write_all_parameters_in_fortran()
        
        # 计算频散谱
        spectrum = tds.calculate_dispersion_spectrum_arrays(
            self.F.flatten(), self.V.flatten()
        )
        spectrum = np.abs(spectrum.reshape(self.obsed_FV.shape))

        if self.DEBUG:
            tds.save_parameters("grt.dat")
            te = time.time()
            print(f"频散谱{spectrum.size}计算完成，耗时: {te-ts:.2f} 秒")


        return spectrum

    def get_dsp_curve(self, vs):

        # 设置媒质参数
        T = 1/self.f_control[-1::-1]
        velocities = calDisp(self.thickness, 
                             self.vp, 
                             vs, 
                             self.rho, 
                             T, 
                wave='rayleigh', mode=1, velocity='phase',
                Qs=self.qs, Qp=self.qp)[-1::-1]
        

        return velocities

    def make_node(self, vs):
        """
        定义操作的输入和输出
        """
        vs = pt.as_tensor_variable(vs)
        return Apply(self, [vs], [pt.dscalar()])
    
    def perform(self, node, inputs, outputs):
        """
        执行实际计算
        """
        vs_val, = inputs
        dv = np.diff(vs_val)
        if np.any(dv < 0):
            outputs[0][0] = np.array(-np.inf)
            return

        if self.METHOD == 'spec':        
            try:

                spectrum = self.get_spectrum(vs_val)

                
                smax = np.max(spectrum)
                spectrum[spectrum > smax/10] = smax/10

                log_likelihood = np.mean(spectrum* self.obsed_FV)
                
                outputs[0][0] = np.array(log_likelihood)
                
            except Exception as e:
                print(f"计算出错: {e}")
                outputs[0][0] = np.array(-np.inf)
        elif self.METHOD == 'curve':
            curve = self.get_dsp_curve(vs_val)

            log_likelihood = -np.sum( (curve-self.v_control)**2)/np.sum(2*self.sigma**2)
            
            outputs[0][0] = np.array(log_likelihood)

def generate_fv_mask(freqs, vs, f_control, v_control, v_win_std) -> np.ndarray:
    '''
    生成频散曲线掩码
    '''
    nf = len(freqs)
    nv = len(vs)
    dv = vs[1]-vs[0]
    assert dv>0
    mask = np.zeros([nf, nv])

    disp_curve_func = interpolate.interp1d(f_control, v_control, kind='linear',
                                        bounds_error=False, fill_value=(0,0))
    V_std_func = interpolate.interp1d(f_control, v_win_std, kind='linear',
                                    bounds_error=False, fill_value=(v_win_std[0], v_win_std[-1]))

    for i, freq in enumerate(freqs):
        if freq<min(f_control) or freq>max(f_control):
            continue
        std_i = V_std_func(freq)
        v_i = disp_curve_func(freq)

        mask_i = np.exp(-(vs-v_i) ** 2 / (2*std_i*std_i))
        mask[i, :] = mask_i
    
    return mask


def main():
    """
    主函数：使用黑箱似然函数执行面波频散反演
    """
    print("=" * 60)
    print("面波频散曲线贝叶斯反演（基于黑箱似然函数）")
    print("=" * 60)
    
    # 1. 生成模拟数据
    print("生成四层模型模拟数据...")
    z = np.array([0, 0.1, 0.2, 0.5, 2])
    thickness = z[1:]-z[:-1]
    vs = np.array([0.4, 0.6, 0.8, 1.0])      # 横波速度 (km/s)
    vp = np.array([1.5, 2.0, 2.3, 2.5])      # 纵波速度 (km/s)
    qs = np.array([500.0, 500.0, 500.0, 500.0])  # 横波品质因子
    qp = np.array([1000.0, 1000.0, 1000.0, 1000.0])  # 纵波品质因子
    rho = np.array([2.0, 2.0, 2.0, 2.0])  # 密度
    
    f1 = np.arange(0.1, 3, 0.2)
    c1 = calDisp(thickness, vp, vs, rho, 1/f1[::-1], 
                 wave='rayleigh', mode=1, Qs=qs, Qp=qp)[::-1]
    f2 = np.arange(1, 2.5, 0.2)
    c2 = calDisp(thickness, vp, vs, rho, 1/f2[::-1], 
                 wave='rayleigh', mode=2, Qs=qs, Qp=qp)[::-1]
    
    freqs = np.linspace(0.1, 3, 40)  # 减少点数以提高计算效率
    vels = np.linspace(0.2, 1, 200)

    E_syn1 = generate_fv_mask(freqs, vels, f1, c1, v_win_std=np.array([0.05]*len(f1)))
    E_syn2 = generate_fv_mask(freqs, vels, f2, c2, v_win_std=np.array([0.05]*len(f2)))
    FV_syn = E_syn1 + E_syn2
    FV_syn[FV_syn>1] = 1
    V,F = np.meshgrid(vels, freqs)

    plt.figure(figsize=(10, 6))
    plt.pcolormesh(F, V, FV_syn)
    plt.plot(f1, c1, 'w-', linewidth=2, label='Mode 1')
    plt.plot(f2, c2, 'w--', linewidth=2, label='Mode 2')
    plt.xlabel('Frequency (Hz)')
    plt.ylabel('Phase Velocity (km/s)')
    plt.title('Synthetic Dispersion Data')
    plt.colorbar(label='Energy')
    plt.legend()
    plt.savefig('figs/FV_syn_blackbox.png', dpi=150, bbox_inches='tight')
    plt.close()

    # 2. 设置反演模型
    print("设置反演模型...")

    N_layers = 4
    Z = np.array([0, 0.1, 0.2, 0.5])
    
    # 创建黑箱似然操作
    likelihood_op = DispersionLikelihoodOp(freqs, vels, Z, FV_syn, DEBUG=False)
    # 先验分布 - 剪切波速度
    V_est = vs
    V_std = np.array([0.4, 0.4, 0.4, 0.4])
    spec1 = likelihood_op.get_spectrum(V_est)
    
    # spec1[spec1<spec1.max()/100]=spec1.max()/100
    plt.close('all')
    plt.figure(figsize=(10, 6))
    
    plt.pcolormesh(F, V, spec1,cmap='nipy_spectral_r')
    plt.plot(f1, c1, 'w-', linewidth=2, label='Mode 1')
    plt.plot(f2, c2, 'w--', linewidth=2, label='Mode 2')
    plt.title(V_est)
    plt.xlabel('Frequency (Hz)')
    plt.ylabel('Phase Velocity (km/s)')
    plt.colorbar()
    plt.savefig(f'figs/test.FV.png', dpi=300, bbox_inches='tight')
    plt.close()

    # assert 1==2
    
    # 定义PyMC模型
    with pm.Model() as model:
        
        
        # 使用Lognormal先验，确保Vs为正值
        vs = pm.Lognormal('vs', mu=np.log(V_est), sigma=V_std, shape=N_layers)
        
        # 定义似然函数
        # log_likelihood = pm.DensityDist(
        #     'likelihood', 
        #     vs, 
        #     logp=likelihood_op, 
        #     observed=0  # 黑箱函数已经包含了观测数据
        # )
        def custom_dist_loglike(data, x):
            # data, or observed is always passed as the first input of CustomDist
            return likelihood_op(x)
        
        
        likelihood = pm.CustomDist(
                "likelihood", vs, observed=np.zeros(N_layers),logp=custom_dist_loglike
                )
        
        # 采样
        print("开始采样...")
        trace = pm.sample(
            draws=1000,             # 采样次数
            tune=500,              # 调谐次数
            # chains=4,               # 链数
            # cores=4,                # 使用的CPU核心数
            # target_accept=0.9,      # 目标接受率
            # nuts_sampler='numpyro'
            # return_inferencedata=True,
            # compute_convergence_checks=True
        )
        
        # 打印结果摘要
        print(az.summary(trace))
        
    
    # 绘制后验预测结果
    posterior_samples = trace.posterior['vs'].values
    n_samples = posterior_samples.shape[0] * posterior_samples.shape[1]
    
    # 获取对数似然值并找到最大值的样本
    log_likelihood_values = trace.sample_stats['lp'].values
    flat_log_likelihood = log_likelihood_values.flatten()
    
    # 获取对数似然最高的几个样本索引
    top_indices = np.argpartition(flat_log_likelihood, -5)[-5:]  # 获取前10个最高似然样本
    
    
    # 绘制对数似然最高的样本对应的频散曲线
    for idx in top_indices:
        chain_idx = idx // log_likelihood_values.shape[1]
        draw_idx = idx % log_likelihood_values.shape[1]
        sample_vs = posterior_samples[chain_idx, draw_idx, :]
        likelihood_value = flat_log_likelihood[idx]
        

        spec1 = likelihood_op.get_spectrum(sample_vs)
        spec1[spec1<spec1.max()/10]=spec1.max()/10

        plt.close('all')
        plt.figure(figsize=(10, 6))
        
        plt.pcolormesh(F, V, spec1,cmap='nipy_spectral_r')
        plt.plot(f1, c1, 'w-', linewidth=2, label='Mode 1')
        plt.plot(f2, c2, 'w--', linewidth=2, label='Mode 2')
        plt.title(likelihood_value)
        plt.xlabel('Frequency (Hz)')
        plt.ylabel('Phase Velocity (km/s)')

        plt.savefig(f'figs/{idx:02d}.png', dpi=300, bbox_inches='tight')
        plt.close()
    
    return trace


if __name__ == "__main__":
    # 运行反演
    trace = main()